#!/bin/bash

read -r -d '' USAGE <<-'USAGE_END'
USAGE:

        vm [ -h | --help ]
        vm [ -l | --list ]
        vm [ <img> | <img-dir> | <n> ] [ <qemu-args> ]

    where: <img> is the path to an .img file
    where: <img-dir> is one of `ls -d $VM_DIR/*/`
    where: <n> is one of the numbered selections of `vm --list`
    where: <qemu-args> are passed to qemu


ENVIRONMENT:

    * VM_DIR       - (mandatory) full path to parent directory of VM subdirectories
    * ARCH         - (optional)  launches qemu-system-$ARCH (default: 'x86_64 + kvm + smp')
    * MEM          - (optional)  specifies VM memory size (default: 2G)
    * AUDIODEV     - (optional)  specifies virtual audo device
    * ISO          - (optional)  full path to .iso file for first optical disc - set as primary boot media (default: none)
    * IMG_B        - (optional)  full path to .img file for second hard disk (default: none)
    * SMP          - (optional)  SMP options as CSV: N_CORES,N_THREADS,N_SOCKETS (default: 2,1,1)
    * VM_SSH_PORT  - (optional)  host port to bind to guest port tcp/22
    * VM_HTTP_PORT - (optional)  host port to bind to guest port tcp/80
    * VM_SSH_LOGIN - (optional)  guest SSH user to login as (if VM_SSH_PORT defined)
    * VM_WAIT_SSH  - (optional)  pause before SSH login (if VM_SSH_PORT and VM_SSH_LOGIN defined)


EXAMPLES:

    vm                     # present an alphabetical list of VMs, and wait for choice
    vm /home/me/vm.img     # launch a VM by path to an image file (absolute or relative)
    vm img-dir             # launch $VM_DIR/img-dir/img-dir.img if it exists
    vm -h (or --help)      # present this message then quit
    vm -l (or --list)      # present an alphabetical list of VMs, then quit
    vm 1                   # launch a VM by it's current order number (per `vm --list`)
    ARCH=i386          vm  # launch 32-bit x86 VM
    MEM=1G             vm  # launch VM using 1GB of system memory
    ISO=/path/to/iso   vm  # boot from .iso file
    IMG_B=/path/to/img vm  # mount second hard disk
    VM_SSH_PORT=       vm  # do not forward SSH
    VM_HTTP_PORT=      vm  # do not forward HTTP
    VM_SSH_LOGIN=      vm  # forward SSH but do not login
    VM_WAIT_SSH=60     vm  # wait N seconds to login
USAGE_END
readonly HAS_LIST_SWITCH=$([ "$1" == "-l" -o "$1" == "--list" ] && echo 1 || echo 0)
readonly HAS_HELP_SWITCH=$([ "$1" == "-h" -o "$1" == "--help" ] && echo 1 || echo 0)
readonly IMG_NAME_OR_N=$1                                          # pending validation
readonly IMG_FULL_PATH="$VM_DIR/$IMG_NAME_OR_N/$IMG_NAME_OR_N.img" # pending validation
readonly QEMU_ARGS=$(args="$*" ; echo $args | grep ' -' > /dev/null && echo "-${args#* -}" )
readonly IS_HEADLESS=$( [[ ! -f /usr/lib/qemu/ui-sdl.so ]] && echo 1 || echo 0 )
readonly VM_DIR_NULL_ERR_MSG="\$VM_DIR must be defined in the environment"
readonly VM_DIR_MISSING_ERR_MSG="\$VM_DIR not found: '$VM_DIR'"
readonly IMG_MOUNTED_ERR_MSG="\$Img is mounted"
readonly ISOS_DIR_NAME='isos'
readonly SCRATCH_IMG_NAME='scratch'
readonly SCRATCH_IMG_SIZE=1G
readonly ISO_NONE='None'
readonly ISOS_DIR=$VM_DIR/$ISOS_DIR_NAME
readonly SCRATCH_IMG_DIR=$VM_DIR/$SCRATCH_IMG_NAME
readonly SCRATCH_IMG=$SCRATCH_IMG_DIR/$SCRATCH_IMG_NAME.img
readonly SMP_REGEX="^[0-9],[0-9],[0-9]$"
readonly DEF_MEM='2G'
readonly DEF_AUDIODEV='alsa'
readonly DEF_VGA='std' # cirrus, qxl, std, vmware
readonly USE_VIRTIO=0
readonly USE_VIRTFS=0
readonly USE_VIRTSMB=0
IMG_DIRS=() # deferred
ISOS=()     # deferred
N_IMGS=0    # deferred
N_ISOS=0    # deferred

ImgDir=        # deferred
OverridesFile= # deferred
Img=           # deferred
Iso=           # deferred


SHOULD_CLR_SCREEN=1
Clear() { (( SHOULD_CLR_SCREEN )) && clear ; }

PopulateImgDirs()
{
  (( ${#IMG_DIRS[@]} == 0 )) || return

  local img_full_path dir img

  IMG_DIRS=( $(
    [[ -d "$SCRATCH_IMG_DIR/" ]] || mkdir -p "$SCRATCH_IMG_DIR"                      >&2
    [[ -f "$SCRATCH_IMG"      ]] || qemu-img create "$SCRATCH_IMG" $SCRATCH_IMG_SIZE >&2
    [[ -f "$SCRATCH_IMG"      ]] && [[ "$IMG_B" != "$SCRATCH_IMG" ]] && echo $SCRATCH_IMG_NAME
    for img_full_path in $(ls -d $VM_DIR/*/*.img 2> /dev/null)
    do  dir=$(basename $(dirname $img_full_path))
        img=$(basename           $img_full_path )
        [[ "$dir" != "$SCRATCH_IMG_NAME" ]] && [[ "$img_full_path" != "$IMG_B" ]] || continue

        grep "^$dir\.img$" <<<$img > /dev/null && echo $dir
    done
  ) )
  readonly IMG_DIRS
  readonly N_IMGS=${#IMG_DIRS[@]}
}

PopulateIsos()
{
  (( ${#ISOS[@]} == 0 )) || return

  local unsorted_isos sorted_isos iso_n sorted_iso_n

  unsorted_isos=( $(
    [[ -d "$ISOS_DIR/" ]] || mkdir -p "$ISOS_DIR"

    find $ISOS_DIR/ -name *.iso 2> /dev/null
  ) )
  sorted_isos=( $(
    for (( iso_n = 0 ; iso_n < ${#unsorted_isos[@]} ; ++iso_n ))
    do  echo "${unsorted_isos[$iso_n]##*\/}/$iso_n"
    done | sort
  ) )
  for (( iso_n = 0 ; iso_n < ${#sorted_isos[@]} ; ++iso_n ))
  do  sorted_iso_n=${sorted_isos[$iso_n]#*\/}

      ISOS[$iso_n]="${unsorted_isos[$sorted_iso_n]}"
  done
  ISOS[$iso_n]=${ISO_NONE}

  readonly ISOS
  readonly N_ISOS=${#ISOS[@]}
}

DoesImgExist() { [[ -f "$1" ]] && [[ "$(grep -E '^.+\.img$' <<<$1)" == "$1" ]] ; }

DoesIsoExist() { [[ -f "$1" ]] && [[ "$(grep -E '^.+\.iso$' <<<$1)" == "$1" ]] ; }

IsImgMounted() { DoesImgExist "$1" && grep "$1" < <(losetup -l) ; return $(( $? )) ; }

IsInteger() { [[ "$1" =~ ^([0-9]+)$ ]] ; }

IsValidSelection() # (option_n , n_options)
{
  local option_n=$( IsInteger $1 && echo $1 || echo '-1')
  local n_options=$(IsInteger $2 && echo $2 || echo '0' )

  [[ "$option_n" -ge "0" ]] && [[ "$option_n" -le "$n_options" ]]
}

PrintImgOptions()
{
  local img_dir_n

  for img_dir_n in "${!IMG_DIRS[@]}" ; do echo "$(( $img_dir_n + 1 )) ${IMG_DIRS[$img_dir_n]}" ; done ;
}

PrintIsoOptions()
{
  local iso_n

  for iso_n in "${!ISOS[@]}" ; do echo "$(( $iso_n + 1 )) ${ISOS[$iso_n]##*/}" ; done ;
}

VmDlg() # (err_msg dialog_args*)
{
  local err_msg=$( [[ -n "$1" ]] && echo -n "\Z1$1\Zn" ) ; shift ;

  dialog --stdout --backtitle "KISS VM Manager" --colors --title "$err_msg" "$@"
}

SelectImage() # (img_n dlg_err_msg)
{
  local img_n=$1
  local dlg_err_msg="$( [[ -z "$ISO" ]] || DoesIsoExist "$ISO" || echo "\$ISO not found")"
  local img_selection=$(IsValidSelection $img_n $N_IMGS && echo $img_n || echo '-1')

  if   (( $N_IMGS > 0 ))
  then IsValidSelection $img_selection $N_IMGS                      || \
       img_selection=$( VmDlg "$dlg_err_msg"                           \
                              --menu "Select a primary disk:" 20 70 50 \
                              $(PrintImgOptions)                       )
       Clear
       [[ "$img_selection" != '' ]] || return 1
  else echo "no conventionally-named images found in \$VM_DIR '$VM_DIR/'"
       return 1
  fi

  ImgDir=${IMG_DIRS[$(( $img_selection - 1 ))]}
  OverridesFile=$VM_DIR/$ImgDir/vm-config
  Img=$VM_DIR/$ImgDir/$ImgDir.img
}

SelectIso()
{
  local iso_selection='-1'

  if (( $N_ISOS > 0 ))
  then IsValidSelection $iso_selection $N_ISOS                     || \
       iso_selection=$( VmDlg '' --menu "Select a boot ISO:" 20 70 50 \
                              $(PrintIsoOptions)                      )
       Clear
       [[ "$iso_selection" != '' ]] || return 1
  else echo "no ISOs found in \$ISOS_DIR: '$ISOS_DIR/'"
       return 1
  fi

  Iso=${ISOS[$(( $iso_selection - 1 ))]}
}

WaitSsh()
{
  ssh-keygen -f "~/.ssh/known_hosts" -R [localhost]:$VM_SSH_PORT
  while (( WAIT_SSH > 0 ))
  do    WAIT_SSH=($WAIT_SSH-1)
        Clear ; echo "logging in ssh in $(($WAIT_SSH+1)) seconds" ;
        sleep 1
  done
}


## main entry ##

(( $HAS_LIST_SWITCH )) && PrintImgOptions | sed 's|^\([0-9]*\) |\1) |' && exit
(( $HAS_HELP_SWITCH )) && echo "$USAGE"                                && exit
[[   -z "$VM_DIR"   ]] && echo "$VM_DIR_NULL_ERR_MSG"                  && exit
[[ ! -d "$VM_DIR"   ]] && echo "$VM_DIR_MISSING_ERR_MSG"               && exit


# initialize data
PopulateImgDirs ; PopulateIsos ;

# determine which VM image to launch
if   [[ -n "$IMG_NAME_OR_N" ]]
then if ! IsInteger $IMG_NAME_OR_N
     then # find image by img name
          for img in "$IMG_NAME_OR_N" "$IMG_FULL_PATH" ; do DoesImgExist $img && Img=$img ; done ;
          [[ "$Img" ]]                            || echo "\$Img not found at '$IMG_NAME_OR_N' or '$IMG_FULL_PATH'"
     else IsValidSelection $IMG_NAME_OR_N $N_IMGS || echo "\$Selection ($IMG_NAME_OR_N) out of range"
     fi
fi

# find image by img_n, or prompt
DoesImgExist $Img || SelectImage "$IMG_NAME_OR_N" || exit 1
IsImgMounted $Img && echo "$IMG_MOUNTED_ERR_MSG"  && exit 1

# determine which ISO to boot
Iso="$ISO"
if   ! DoesIsoExist "$Iso" && [[ "$ImgDir" == "$SCRATCH_IMG_NAME" ]]
then (( $(ls *.iso | wc -l) == 1 )) && Iso=$(ls *.iso) || true # pwd override

     until DoesIsoExist "$Iso" || [[ "$Iso" == "${ISO_NONE}" ]] ; do SelectIso || exit 1 ; done ;
fi

# log results
echo -e "Selected:$( [[ "$Iso" ]] && echo "\n\tIso: $Iso")\n\tImg: $Img"

# prepare the environment
[[ -f $OverridesFile    ]] && source $OverridesFile
[[ "$SMP" =~ $SMP_REGEX ]] && N_CORES=${BASH_REMATCH[1]} N_THREADS=${BASH_REMATCH[2]} N_SOCKETS=${BASH_REMATCH[3]} || \
                              N_CORES=2                  N_THREADS=1                  N_SOCKETS=1
[[ -z "$ARCH"           ]] && ARCH="x86_64"
[[ "$ARCH" == 'x86_64'  ]] && CPU="-cpu host -smp cores=$N_CORES,threads=$N_THREADS,sockets=$N_SOCKETS"
[[ "$ARCH" == 'x86_64'  ]] && ARCH="x86_64 -enable-kvm"
[[ "$ImgDir"            ]] && NAME="-name $ImgDir"
[[ "$IMG_B"             ]] && HD="-drive file=$IMG_B,format=raw,cache=writeback"
DoesIsoExist "$Iso"        && CD="-cdrom $Iso" BOOT="-boot order=d"                 || CD="" BOOT=""
[[ "$MEM"               ]] && MEM="-m $MEM"                                         || MEM="-m $DEF_MEM"
[[ "$VGA"               ]] && VGA="-vga $VGA"                                       || VGA="-vga $DEF_VGA"
[[ "$AUDIODEV"          ]] && AUDIODEV="-audiodev $AUDIODEV"                        || AUDIODEV="-audiodev $DEF_AUDIODEV"
[[ "$VM_SSH_PORT"       ]] && FWD_SSH=",hostfwd=tcp::$VM_SSH_PORT-:22"              || FWD_SSH=''
[[ "$VM_HTTP_PORT"      ]] && FWD_HTTP=",hostfwd=tcp::$VM_HTTP_PORT-:$VM_HTTP_PORT" || FWD_HTTP=''
QEMU="qemu-system-$ARCH $NAME"
HD="-drive file=$Img,format=raw,cache=writeback $HD"
# BOOT="-boot menu=on"

# prepare networking
SHARED_DIR_DEFAULT=$HOME/.config/vm-shared
[[ ! -d "$SHARED_DIR" ]] && [[ -d "$SHARED_DIR_DEFAULT" ]] && SHARED_DIR=$SHARED_DIR_DEFAULT
NET_VIRTIO="-netdev user,id=vmnic$FWD_SSH$FWD_HTTP -device virtio-net,netdev=vmnic"
NET_VIRTSMB="-net user,smb=\"$SHARED_DIR\" -net nic,model=virtio"
NET_VIRTFS="-virtfs local,path=\"$SHARED_DIR\",mount_tag=host0,security_model=passthrough,id=host0"
# SSH="-redir tcp:$VM_SSH_PORT::22" # no virtio - deprecated
if   [[ "$VM_SSH_PORT" ]] || [[ "$VM_HTTP_PORT"    ]]
then [[ "$VM_SSH_PORT" ]] && [[ "$VM_SSH_LOGIN"    ]] && SSH="ssh -o StrictHostKeyChecking=no -p $VM_SSH_PORT $VM_SSH_LOGIN@localhost"
     [[ "`echo $VM_WAIT_SSH | grep -E '^[1-9]+$'`" ]] || WAIT_SSH=30
     NET=$( ( (( $USE_VIRTIO  )) && echo "$NET_VIRTIO"  )
            ( (( $USE_VIRTSMB )) && echo "$NET_VIRTSMB" )
            ( (( $USE_VIRTFS  )) && echo "$NET_VIRTFS"  ) ) # guest fstab entry: host0 /a-mountpoint 9p trans=virtio,version=9p2000.L 0 0
fi

# prepare A/V
# AUDIO='-device intel-hda -device hda-duplex'
AUDIO="${AUDIODEV},id=myaudiodev -device AC97,audiodev=myaudiodev"
VIDEO_VNC='-display vnc=:0'
VIDEO_SDL='-display sdl' # ,show-cursor=on -no-frame
(( ${IS_HEADLESS} )) && VIDEO="${VIDEO_VNC}" || VIDEO="${VGA} ${VIDEO_SDL}"

MISC=''


# load virtio kernel modules
if   (( $USE_VIRTIO ))
then virtio_modules=`lsmod | grep virtio`
     [[ "`echo $virtio_modules | grep virtio_net`"    == "" ]] && su -c 'modprobe virtio-net'
     #[[ "`echo $virtio_modules | grep virtio_blk`"     == "" ]] && su -c 'modprobe virtio-blk'
     #[[ "`echo $virtio_modules | grep virtio_scsi`"    == "" ]] && su -c 'modprobe virtio-scsi'
     #[[ "`echo $virtio_modules | grep virtio_serial`"  == "" ]] && su -c 'modprobe virtio-serial'
     #[[ "`echo $virtio_modules | grep virtio_balloon`" == "" ]] && su -c 'modprobe virtio-balloon'
fi


# concatenate command line and report
[[ "$FWD_SSH"  ]] && ( [[ "$NET" ]] && echo "SSH port 22 forwarded to localhost:$VM_SSH_PORT" || \
                                       echo "no USE_VIRT* enabled - SSH will not be forwarded"   )
[[ "$FWD_HTTP" ]] && ( [[ "$NET" ]] && echo "HTTP port $VM_HTTP_PORT forwarded to localhost:$VM_HTTP_PORT" || \
                                       echo "no USE_VIRT* enabled - HTTP will not be forwarded"               )
CMD="$QEMU $CPU $MEM $HD $CD $BOOT $NET $AUDIO $VIDEO $MISC $QEMU_ARGS"
msg="launching vm: '$(echo "$Img" | grep -E '^.+\.img$' | sed -e 's/^\(.*\/\)\?\(.*\).img$/\2/')'"


# launch VM and optionally login ssh
if   [[ "$1" == '--cmd' ]]
then echo $CMD
elif [[ "$SSH" ]] && [[ "$WAIT_SSH" ]]
then echo $msg ; $CMD & WaitSsh && $SSH ;
else echo $msg ; $CMD ;
fi
